# This script will read the weights from a baseline DNN model of hidden layer size 1124 and create a 
# Strutured DNN model with standard weights of 1024 and speaker weights of 100
#

import numpy as np
import json
from StringIO import StringIO

def array_2_string(array):
  str_out = StringIO()
  np.savetxt(str_out, array)
  return str_out.getvalue()

def string_2_array(string):
  str_in = StringIO(string)
  return np.loadtxt(str_in)

def main():
  nnet_dict = {}
  num_layers=6
  hlayer_size=1024
  slayer_size=100
  ivec_size=25
  base_dnn='exp_pdnn/warm6_str_dnn_fmllr/dnn.temp'
  ptr_file='exp_pdnn/warm6_str_dnn_fmllr/dnn.ptr'
  base_dnn_file = open(base_dnn,'r')
  nnet_dict = json.load(base_dnn_file)

  rng = np.random.RandomState(89677)
  W20_values = np.asarray(rng.uniform(low=-np.sqrt(6. / (ivec_size + hlayer_size)),high=np.sqrt(6. / (ivec_size + hlayer_size)),size=(ivec_size, hlayer_size)))
  W30_values = np.asarray(rng.uniform(low=-np.sqrt(6. / (ivec_size + hlayer_size)),high=np.sqrt(6. / (ivec_size + hlayer_size)),size=(ivec_size, slayer_size)))
  #W2_values = np.asarray(rng.uniform(low=-np.sqrt(6. / (100 + 1024)),high=np.sqrt(6. / (100 + 1024)),size=(100, 1024)))
  #W3_values = np.asarray(rng.uniform(low=-np.sqrt(6. / (100 + 1024)),high=np.sqrt(6. / (100 + 1024)),size=(100, 100)))
  #b2_values = np.zeros((100))

  #ogregw_key = 'logreg W'
  #logW = string_2_array(nnet_dict[logregw_key])
  #temp = np.asarray(rng.uniform(low=-np.sqrt(6. / (440 + 1024)),high=np.sqrt(6. / (440 + 1024)),size=(100, 2042)))
  #nnet_dict[logregw_key]=array_2_string(np.vstack((logW,temp)))

  for i in range(num_layers):
    key_w1 = str(i) + ' sigmoid W1'
    key_w2 = str(i) + ' sigmoid W2'
    key_w3 = str(i) + ' sigmoid W3'
    key_b1 = str(i) + ' sigmoid b1'
    key_b2 = str(i) + ' sigmoid b2'

    W=string_2_array(nnet_dict.pop(str(i) + ' sigmoid W'))
    b=string_2_array(nnet_dict.pop(str(i) + ' sigmoid b'))

    nnet_dict[key_b1]=array_2_string(b[:hlayer_size])
    nnet_dict[key_b2]=array_2_string(b[hlayer_size:])
    if i == 0:
      nnet_dict[key_w1]=array_2_string(W[:,:hlayer_size])
      nnet_dict[key_w2]=array_2_string(W20_values)
      nnet_dict[key_w3]=array_2_string(W30_values)
    else:
      nnet_dict[key_w1]=array_2_string(W[:hlayer_size,:hlayer_size])
      nnet_dict[key_w2]=array_2_string(W[hlayer_size:,:hlayer_size])
      nnet_dict[key_w3]=array_2_string(W[hlayer_size:,hlayer_size:])

  with open(ptr_file, 'wb') as fp:
    json.dump(nnet_dict, fp, indent=2, sort_keys = True)
    fp.flush()


if __name__=="__main__":
    main()

